import numpy as np
import torch
import torch.nn as nn


class DNN(nn.Module):

    def __init__(self):
        super(DNN, self).__init__()

        self.model = nn.Sequential(
            nn.Linear(100*480*848*4, 20000),
            nn.LeakyReLU(inplace=True),
            nn.Linear(20000, 200),
            nn.LeakyReLU(inplace=True),
            nn.Linear(200, 10),
            nn.LeakyReLU(inplace=True)
        )

    def forward(self, x):
        x = self.model(x)
        return x


if __name__ == '__main__':
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    data = torch.rand(1,4,480, 848)
    data = data.to(device)
    model = DNN(in_channel=4).to(device)
    out = model(data)
    print(out.shape)